#include "imu.h"
#include "math.h"
#include "bmi088.h"
#include "Quaternion.h"
#include "can.h"
#include "interface.h"
#include "kalman.h"
#include "tim.h"

IMU_TypeDef IMU_Data;
Quaternion_TypeDef IMU_Quaternion_Data;

int cal_time_2=0;

// 新增：从加速度计计算姿态角
void IMU_Get_Angle_From_Acc(IMU_TypeDef* m, float ax, float ay, float az, float* roll_acc, float* pitch_acc)
{
    // 归一化加速度
    float norm = sqrtf(ax*ax + ay*ay + az*az);
    if (norm > 1e-6f) {
        ax /= norm;
        ay /= norm;
        az /= norm;
    }
    
    // 从加速度计计算roll和pitch（弧度）
    *roll_acc = atan2f(ay, az);
    *pitch_acc = atan2f(-ax, sqrtf(ay*ay + az*az));
}

// 新增：从四元数获取姿态角
void IMU_Get_Angle_From_Quaternion(IMU_TypeDef* m, float* roll_gyro, float* pitch_gyro, float* yaw_gyro)
{
    float q0 = m->q0, q1 = m->q1, q2 = m->q2, q3 = m->q3;
    
    // roll (x-axis rotation)
    float sinr_cosp = 2.0f * (q0 * q1 + q2 * q3);
    float cosr_cosp = 1.0f - 2.0f * (q1 * q1 + q2 * q2);
    *roll_gyro = atan2f(sinr_cosp, cosr_cosp);
    
    // pitch (y-axis rotation)
    float sinp = 2.0f * (q0 * q2 - q3 * q1);
    if (fabsf(sinp) >= 1)
        *pitch_gyro = copysignf(3.14159265358979323846f / 2.0f, sinp);
    else
        *pitch_gyro = asinf(sinp);
    
    // yaw (z-axis rotation)
    float siny_cosp = 2.0f * (q0 * q3 + q1 * q2);
    float cosy_cosp = 1.0f - 2.0f * (q2 * q2 + q3 * q3);
    *yaw_gyro = atan2f(siny_cosp, cosy_cosp);
}

void IMU_Data_Init(IMU_TypeDef* imu, float n, float delta_time)
{
    imu->n = n;
    imu->IMU_Filiter[0].n = n;
    imu->IMU_Filiter[1].n = n;
    imu->IMU_Filiter[2].n = n;
    imu->delta_time = delta_time;
    imu->q0 = 1.0;  // 初始化四元数为单位四元数
    imu->q1 = 0;
    imu->q2 = 0;
    imu->q3 = 0;
    
    // 初始化Kalman滤波器 - 现在用于角度融合
    imu->IMU_Kalman[0].error_cov = 1;
    imu->IMU_Kalman[0].estimate = 0.1;
    imu->IMU_Kalman[0].DT = delta_time;
    imu->IMU_Kalman[0].process_noise = 0.01f;   // 过程噪声 - 陀螺仪信任度
    imu->IMU_Kalman[0].measure_noise = 0.1f;    // 测量噪声 - 加速度计信任度
    
    imu->IMU_Kalman[1].error_cov = 1;
    imu->IMU_Kalman[1].estimate = 0.1;
    imu->IMU_Kalman[1].DT = delta_time;
    imu->IMU_Kalman[1].process_noise = 0.01f;
    imu->IMU_Kalman[1].measure_noise = 0.1f;
    
    imu->IMU_Kalman[2].error_cov = 1;
    imu->IMU_Kalman[2].estimate = 0.1;
    imu->IMU_Kalman[2].DT = delta_time;
    imu->IMU_Kalman[2].process_noise = 0.01f;
    imu->IMU_Kalman[2].measure_noise = 10.0f;   // yaw的测量噪声很大，因为加速度计无法测量yaw
}

void IMU_Error_Cal(IMU_Filiter_TypeDef* m, float x)
{
    if (m->cnt == m->n)
    {
        m->average = m->sum / m->n;
        m->flag = 1;
        return;
    }
    m->sum += x;
    m->cnt++;
}

float IMU_bessel_filter(IMU_Last_TypeDef* m, float now, float filter)
{
    // 计算滤波后的值
    float temp = 1 * (1 - filter) * (1 - filter) * (1 - filter) * (1 - filter) * (1 - filter) * now +
        5 * (1 - filter) * (1 - filter) * (1 - filter) * (1 - filter) * filter * m->last_0 +
        10 * (1 - filter) * (1 - filter) * (1 - filter) * filter * filter * m->last_1 +
        10 * (1 - filter) * (1 - filter) * filter * filter * filter * m->last_2 +
        5 * (1 - filter) * filter * filter * filter * filter * m->last_3 +
        1 * filter * filter * filter * filter * filter * m->last_4;

    // 更新历史值
    m->last_4 = m->last_3;
    m->last_3 = m->last_2;
    m->last_2 = m->last_1;
    m->last_1 = m->last_0;
    m->last_0 = now;
    return temp;
}

void IMU_update(IMU_TypeDef* m, float gx, float gy, float gz, float ax, float ay, float az)
{
    m->gx = IMU_bessel_filter(&m->Last[0], gx, 0.12);
    m->gy = IMU_bessel_filter(&m->Last[1], gy, 0.12);
    m->gz = IMU_bessel_filter(&m->Last[2], gz, 0.25);
    m->ax = IMU_bessel_filter(&m->Last[4], ax, 0.12);
    m->ay = IMU_bessel_filter(&m->Last[5], ay, 0.12);
    m->az = IMU_bessel_filter(&m->Last[6], az, 0.12);
}

// 新增：从欧拉角转换为四元数
void IMU_Euler_to_Quaternion(float roll, float pitch, float yaw, float* q0, float* q1, float* q2, float* q3)
{
    float cr = cosf(roll * 0.5f);
    float sr = sinf(roll * 0.5f);
    float cp = cosf(pitch * 0.5f);
    float sp = sinf(pitch * 0.5f);
    float cy = cosf(yaw * 0.5f);
    float sy = sinf(yaw * 0.5f);

    *q0 = cr * cp * cy + sr * sp * sy;
    *q1 = sr * cp * cy - cr * sp * sy;
    *q2 = cr * sp * cy + sr * cp * sy;
    *q3 = cr * cp * sy - sr * sp * cy;
}

void IMU_Cal(IMU_TypeDef* m)
{
    float gx = m->gx;
    float gy = m->gy;
    float gz = m->gz;
    float ax = m->ax, ay = m->ay, az = m->az;
    float q0 = m->q0, q1 = m->q1, q2 = m->q2, q3 = m->q3;

    // 四元数归一化
    float norm = sqrtf(q0*q0 + q1*q1 + q2*q2 + q3*q3);
    if (norm > 1e-12f) {
        q0 /= norm;
        q1 /= norm;
        q2 /= norm;
        q3 /= norm;
    }

    // ========== 角度融合核心部分 ==========
    
    // 步骤1：从陀螺仪积分得到预测角度
    float roll_gyro, pitch_gyro, yaw_gyro;
    IMU_Get_Angle_From_Quaternion(m, &roll_gyro, &pitch_gyro, &yaw_gyro);
    
    // 步骤2：从加速度计得到测量角度
    float roll_acc, pitch_acc;
    IMU_Get_Angle_From_Acc(m, ax, ay, az, &roll_acc, &pitch_acc);
    
    // 步骤3：使用Kalman滤波器进行角度融合
    // Roll轴融合
    Kalman_predict(&m->IMU_Kalman[0], roll_gyro);
    float roll_fused = Kalman_update(&m->IMU_Kalman[0], roll_acc);
    
    // Pitch轴融合  
    Kalman_predict(&m->IMU_Kalman[1], pitch_gyro);
    float pitch_fused = Kalman_update(&m->IMU_Kalman[1], pitch_acc);
    
//    // Yaw轴 - 只用陀螺仪积分（加速度计无法测量yaw）
//    Kalman_predict(&m->IMU_Kalman[2], yaw_gyro);
//    // 由于没有磁力计，yaw只能依赖陀螺仪积分
//    float yaw_fused = m->IMU_Kalman[2].estimate;
    float yaw_fused = yaw_gyro;
    // 步骤4：将融合后的角度转换回四元数
    IMU_Euler_to_Quaternion(roll_fused, pitch_fused, yaw_fused, &q0, &q1, &q2, &q3);
    
    // 步骤5：用融合后的四元数作为基础，再用陀螺仪进行短期预测
    // 这样可以保持动态响应性能
    float dt2 = m->delta_time * 0.5f;
    float tq0 = q0, tq1 = q1, tq2 = q2, tq3 = q3;
    
    // 限制角速度
    const float MAX_RATE = 20.0f;
    gx = (fabsf(gx) > MAX_RATE) ? copysignf(MAX_RATE, gx) : gx;
    gy = (fabsf(gy) > MAX_RATE) ? copysignf(MAX_RATE, gy) : gy;
    gz = (fabsf(gz) > MAX_RATE) ? copysignf(MAX_RATE, gz) : gz;
    
    // 用陀螺仪更新四元数（短期预测）
    q0 = tq0 + (-tq1*gx - tq2*gy - tq3*gz) * dt2;
    q1 = tq1 + ( tq0*gx + tq2*gz - tq3*gy) * dt2;
    q2 = tq2 + ( tq0*gy - tq1*gz + tq3*gx) * dt2;
    q3 = tq3 + ( tq0*gz + tq1*gy - tq2*gx) * dt2;
    
    // 最终归一化
    norm = sqrtf(q0*q0 + q1*q1 + q2*q2 + q3*q3);
    if (norm > 1e-12f) {
        q0 /= norm;
        q1 /= norm;
        q2 /= norm;
        q3 /= norm;
    } else {
        q0 = 1.0f; q1 = q2 = q3 = 0.0f;
    }

    m->q0 = q0; m->q1 = q1; m->q2 = q2; m->q3 = q3;
    
    // 保存融合后的角度用于调试
    m->roll_fused = roll_fused * 57.2957795f;
    m->pitch_fused = pitch_fused * 57.2957795f;
    m->yaw_fused = yaw_fused * 57.2957795f;
}

void IMU_to_euler(IMU_TypeDef* m)
{
    float q0 = m->q0;
    float q1 = m->q1;
    float q2 = m->q2;
    float q3 = m->q3;

    // 防止asin溢出
    float sinp = 2.0f * (q0 * q2 - q3 * q1);
    if (fabs(sinp) >= 1)
        m->pitch = copysignf(90.0f, sinp);
    else
        m->pitch = asinf(sinp) * 57.2957795f;

    m->roll = atan2f(2.0f * (q0 * q1 + q2 * q3),
        1.0f - 2.0f * (q1 * q1 + q2 * q2)) * 57.2957795f;

    m->yaw = atan2f(2.0f * (q0 * q3 + q1 * q2),
        1.0f - 2.0f * (q2 * q2 + q3 * q3)) * 57.2957795f;
}

void IMU_exe()
{
    if (IMU_Data.IMU_Filiter[0].flag == 0 ||
        IMU_Data.IMU_Filiter[1].flag == 0 ||
        IMU_Data.IMU_Filiter[2].flag == 0)
    {
        IMU_Error_Cal(&IMU_Data.IMU_Filiter[0], bmi088_data.gyro_rate.x);
        IMU_Error_Cal(&IMU_Data.IMU_Filiter[1], bmi088_data.gyro_rate.y);
        IMU_Error_Cal(&IMU_Data.IMU_Filiter[2], bmi088_data.gyro_rate.z);
    }
    else
    {
        IMU_update(&IMU_Data,
            bmi088_data.gyro_rate.x,
            bmi088_data.gyro_rate.y,
            bmi088_data.gyro_rate.z - IMU_Data.IMU_Filiter[2].average,
            bmi088_data.acc.x,
            bmi088_data.acc.y,
            bmi088_data.acc.z);
        
        __HAL_TIM_SET_COUNTER(&htim1, 0);     
        __HAL_TIM_ENABLE(&htim1);
        
        IMU_Cal(&IMU_Data);
        IMU_to_euler(&IMU_Data);
        
        cal_time_2 = __HAL_TIM_GET_COUNTER(&htim1);
        __HAL_TIM_DISABLE(&htim1);
    }
}